# -*- coding: utf-8 -*-
"""
===============================
@Author     : Zuo WenTao
@Time       : 2024/5/30 9:49
@Description: 
@Change     : 
@File       : monitored_func.py
===============================
"""
from multiprocessing import Process, Queue


class FunctionTimeoutError(Exception):
    pass


def func_wrapper(func, q, *args, **kwargs):
    try:
        result = func(*args, **kwargs)
        q.put(result)
    except Exception as e:
        q.put(e)


def monitored_func(func, timeout):
    def wrapper(*args, **kwargs):
        q = Queue()
        p = Process(target=func_wrapper, args=(func, q, *args), kwargs=kwargs)
        p.start()
        p.join(timeout)

        if p.is_alive():
            p.terminate()
            p.join()
            raise FunctionTimeoutError(f"函数 '{func.__name__}' 超过了执行的超时时间 {timeout} 秒")

        result = q.get()
        if isinstance(result, Exception):
            raise result

        return result

    return wrapper
